import pandas as pd
import numpy as np

# 生成模拟数据
# np.random.seed(42)
# data = []
#
# # 创建3个款号的模拟数据
# for spu in ['A001', 'A002', 'B001']:
#     cart_total = np.random.randint(500, 2000)
#     new_sale = np.random.randint(20, 100)
#
#     # 生成30天的动态数据
#     for day in range(1, 31):
#         base_sale = 100 * np.exp(-0.05 * day)
#         sale_qty = max(0, int(base_sale + np.random.normal(0, 10)))
#
#         data.append({
#             'spu_code': spu,
#             'sum_cart_total': cart_total,
#             'new_sale_qty': new_sale,
#             'days_since_launch': day,
#             'sale_qty': sale_qty
#         })


# 创建DataFrame并保存
df = pd.read_excel("enhanced_sales_data.xlsx")
# df.to_excel("sales_data.xlsx", index=False)
# print("模拟数据已生成并保存为 sales_data.xlsx")

import pandas as pd
import joblib
from sklearn.ensemble import RandomForestRegressor


def train_models(data_file):
    """训练所有款号的模型并保存"""
    # 加载数据
    df = pd.read_excel(data_file)

    # 存储所有训练好的模型
    models = {}

    # 按款号分组训练
    for spu, group in df.groupby('spu_code'):
        # 准备训练数据
        X = group[['sum_cart_total', 'new_sale_qty', 'days_since_launch']]
        y = group['sale_qty']

        # 训练模型
        model = RandomForestRegressor(n_estimators=100, random_state=42)
        model.fit(X, y)

        # 存储模型
        models[spu] = model
        print(f"款号 {spu} 模型训练完成，使用数据量: {len(group)}")

    # 保存所有模型
    joblib.dump(models, 'sales_models.pkl')
    print("所有模型已保存为 sales_models.pkl")

    return models


if __name__ == "__main__":
    models = train_models("enhanced_sales_data.xlsx")
